package codec.huffman;

import util.ColorTransformer;
import util.Log;
import util.Transformer;
import image.Image;
import image.header.JPEGHeader;

public class JPEGCustomizedHuffmanTable  extends JPEGHuffmanTable {
	private Image image;
	
	private JPEGHeader header;
	
	public static final String TAG = "JPEGCustomizedHuffmanTable";
	
	public int[] LUMINANCE_DC_FREQ;
	public int[] LUMINANCE_AC_FREQ;
	public int[] CHROMINANCE_DC_FREQ;
	public int[] CHROMINANCE_AC_FREQ;
	
	public int[] LUMINANCE_DC_CODESIZE;
	public int[] LUMINANCE_AC_CODESIZE;
	public int[] CHROMINANCE_DC_CODESIZE;
	public int[] CHROMINANCE_AC_CODESIZE;
	
	public int[] LUMINANCE_DC_BITS;
	public int[] LUMINANCE_AC_BITS;
	public int[] CHROMINANCE_DC_BITS;
	public int[] CHROMINANCE_AC_BITS;
	
	public int[] LUMINANCE_DC_HUFVAL;
	public int[] LUMINANCE_AC_HUFVAL;
	public int[] CHROMINANCE_DC_HUFVAL;
	public int[] CHROMINANCE_AC_HUFVAL;
	
	public int[] LUMINANCE_DC_HUFSIZE;
	public int[] LUMINANCE_AC_HUFSIZE;
	public int[] CHROMINANCE_DC_HUFSIZE;
	public int[] CHROMINANCE_AC_HUFSIZE;
	
	public int[] LUMINANCE_DC_HUFCODE;
	public int[] LUMINANCE_AC_HUFCODE;
	public int[] CHROMINANCE_DC_HUFCODE;
	public int[] CHROMINANCE_AC_HUFCODE;
	
	public int[] LUMINANCE_DC_EHUFCODE;
	public int[] LUMINANCE_AC_EHUFCODE;
	public int[] CHROMINANCE_DC_EHUFCODE;
	public int[] CHROMINANCE_AC_EHUFCODE;
	
	public int[] LUMINANCE_DC_EHUFSIZE;
	public int[] LUMINANCE_AC_EHUFSIZE;
	public int[] CHROMINANCE_DC_EHUFSIZE;
	public int[] CHROMINANCE_AC_EHUFSIZE;
	
	public JPEGCustomizedHuffmanTable(Image image, JPEGHeader header) {
		this.setImage(image);
		this.setHeader(header);
		init();
	}
	
	private void init() {
		/*
		 * Make FREQ sequence
		 */
		int[][][] freq = generateFrequency(image);
		LUMINANCE_DC_FREQ 		= freq[0][0];
		LUMINANCE_AC_FREQ 		= freq[0][1];
		CHROMINANCE_DC_FREQ 	= freq[1][0];
		CHROMINANCE_AC_FREQ 	= freq[1][1];
		Log.debugArray(TAG, "LUM_DC_FREQ", LUMINANCE_DC_FREQ, 11);
		
		/*
		 * Now that we have the frequencies, starts making codesize
		 */
		LUMINANCE_DC_CODESIZE 		= generateCodeSize(LUMINANCE_DC_FREQ);
		LUMINANCE_AC_CODESIZE 		= generateCodeSize(LUMINANCE_AC_FREQ);
		CHROMINANCE_DC_CODESIZE 	= generateCodeSize(CHROMINANCE_DC_FREQ);
		CHROMINANCE_AC_CODESIZE 	= generateCodeSize(CHROMINANCE_AC_FREQ);
		Log.debugArray(TAG, "LUM_DC_CODESIZE", LUMINANCE_DC_CODESIZE, 11);
		
		/*
		 * Make BITS
		 */
		LUMINANCE_DC_BITS = makeBITS(LUMINANCE_DC_CODESIZE);
		LUMINANCE_AC_BITS = makeBITS(LUMINANCE_AC_CODESIZE);
		CHROMINANCE_DC_BITS = makeBITS(CHROMINANCE_DC_CODESIZE);
		CHROMINANCE_AC_BITS = makeBITS(CHROMINANCE_AC_CODESIZE);
		Log.debugArray(TAG, "LUM_DC_BITS", LUMINANCE_DC_BITS, 11);

		/*
		 * Make HUFVAL sequence
		 */
		LUMINANCE_DC_HUFVAL = makeHUFVAL(LUMINANCE_DC_BITS, LUMINANCE_DC_CODESIZE);
		LUMINANCE_AC_HUFVAL = makeHUFVAL(LUMINANCE_AC_BITS, LUMINANCE_AC_CODESIZE);
		CHROMINANCE_DC_HUFVAL = makeHUFVAL(CHROMINANCE_DC_BITS, CHROMINANCE_DC_CODESIZE);
		CHROMINANCE_AC_HUFVAL = makeHUFVAL(CHROMINANCE_AC_BITS, CHROMINANCE_AC_CODESIZE);
		Log.debugArray(TAG, "LUM_DC_HUFVAL", LUMINANCE_DC_HUFVAL, 11);
		
		/*
		 * Make HUFSIZE sequence
		 */
		LUMINANCE_DC_HUFSIZE = makeHUFSIZE(LUMINANCE_DC_CODESIZE, LUMINANCE_DC_HUFVAL);
		LUMINANCE_AC_HUFSIZE = makeHUFSIZE(LUMINANCE_AC_CODESIZE, LUMINANCE_AC_HUFVAL);
		CHROMINANCE_DC_HUFSIZE = makeHUFSIZE(CHROMINANCE_DC_CODESIZE, CHROMINANCE_DC_HUFVAL);
		CHROMINANCE_AC_HUFSIZE = makeHUFSIZE(CHROMINANCE_AC_CODESIZE, CHROMINANCE_AC_HUFVAL);
		Log.debugArray(TAG, "LUM_DC_HUFSIZE", LUMINANCE_DC_HUFSIZE, 11);
		
		/*
		 * Make HUFCODE sequence
		 */
		LUMINANCE_DC_HUFCODE = generateCode(LUMINANCE_DC_HUFVAL, LUMINANCE_DC_HUFSIZE);
		LUMINANCE_AC_HUFCODE = generateCode(LUMINANCE_AC_HUFVAL, LUMINANCE_AC_HUFSIZE);
		CHROMINANCE_DC_HUFCODE = generateCode(CHROMINANCE_DC_HUFVAL, CHROMINANCE_DC_HUFSIZE);
		CHROMINANCE_AC_HUFCODE = generateCode(CHROMINANCE_AC_HUFVAL, CHROMINANCE_AC_HUFSIZE);
		Log.debugArray(TAG, "LUM_DC_HUFCODE", LUMINANCE_DC_HUFSIZE, 11);
		
		/**
		 * Make EHUF
		 */
		// LUMINANCE DC's ehuf
		int hufval_max = LUMINANCE_DC_HUFVAL[0];
		for (int i = 1; i < LUMINANCE_DC_HUFVAL.length; i++) {
			if (LUMINANCE_DC_HUFVAL[i] > hufval_max) hufval_max = LUMINANCE_DC_HUFVAL[i]; 
		}
		int[][] ehuf = makeEHUF(LUMINANCE_DC_HUFVAL, LUMINANCE_DC_HUFCODE, LUMINANCE_DC_HUFSIZE, hufval_max); 
		LUMINANCE_DC_EHUFCODE = ehuf[0];
		LUMINANCE_DC_EHUFSIZE = ehuf[1];

		// Luminance ac's ehuf
		hufval_max = LUMINANCE_AC_HUFVAL[0];
		for (int i = 1; i < LUMINANCE_AC_HUFVAL.length; i++) {
			if (LUMINANCE_AC_HUFVAL[i] > hufval_max) hufval_max = LUMINANCE_AC_HUFVAL[i]; 
		}
		ehuf = makeEHUF(LUMINANCE_AC_HUFVAL, LUMINANCE_AC_HUFCODE, LUMINANCE_AC_HUFSIZE, hufval_max); 
		LUMINANCE_AC_EHUFCODE = ehuf[0];
		LUMINANCE_AC_EHUFSIZE = ehuf[1];
		
		// Chrominance dc's ehuf
		hufval_max = CHROMINANCE_DC_HUFVAL[0];
		for (int i = 1; i < CHROMINANCE_DC_HUFVAL.length; i++) {
			if (CHROMINANCE_DC_HUFVAL[i] > hufval_max) hufval_max = CHROMINANCE_DC_HUFVAL[i]; 
		}
		ehuf = makeEHUF(CHROMINANCE_DC_HUFVAL, 
				CHROMINANCE_DC_HUFCODE, 
				CHROMINANCE_DC_HUFSIZE, hufval_max); 
		CHROMINANCE_DC_EHUFCODE = ehuf[0];
		CHROMINANCE_DC_EHUFSIZE = ehuf[1];
		
		// chrominance ac's ehuf
		hufval_max = CHROMINANCE_AC_HUFVAL[0];
		for (int i = 1; i < CHROMINANCE_AC_HUFVAL.length; i++) {
			if (CHROMINANCE_AC_HUFVAL[i] > hufval_max) hufval_max = CHROMINANCE_AC_HUFVAL[i]; 
		}
		ehuf = makeEHUF(CHROMINANCE_AC_HUFVAL, 
				CHROMINANCE_AC_HUFCODE, 
				CHROMINANCE_AC_HUFSIZE, hufval_max); 
		CHROMINANCE_AC_EHUFCODE = ehuf[0];
		CHROMINANCE_AC_EHUFSIZE = ehuf[1];
		
		Log.debugArray(TAG, "LUM_DC_EHUFCODE", LUMINANCE_DC_EHUFCODE, 11);
	}
	
	/**
	 * Make hufsize table
	 * @param bits
	 * @param hufval
	 * @return
	 */
	protected int[] makeHUFSIZE(int[] codesize, int[] hufval) {
		int length = hufval.length;
		int[] hufsize = new int[length];
		
		for (int i = 0; i < length; i++) {
			int word = hufval[i];
			hufsize[i] = codesize[word];
		}
		
		return hufsize;
	}
	
	/**
	 * Generate the BITS sequence for the hufman table
	 * @param codesize codesize sequence of this table.
	 * @return BITS sequence
	 */
	protected int[] makeBITS(int[] codesize) {
		int max_clen = 32;
		int bits[] = new int[max_clen + 1];
		/*
	  	 * Finished with codesize, make BITS sequence
	  	 */
	  	for (int i = 0; i <= 256; i++) {
	  		if (codesize[i] > 0) {
	  			bits[codesize[i] - 1]++;
	  		}
	  	}
	  	
	  	/*
	  	 * Make it so there's no code longer than 16
	  	 */
	  	int i;
	  	for (i = max_clen; i > 16; i--) {
	  		while (bits[i] > 0) {
	  			int j = i - 2;		/* find length of new prefix to be used */
	  			while (bits[j] == 0)
	  				j--;
	  	      
	  			bits[i] -= 2;		/* remove two symbols */
	  			bits[i-1]++;		/* one goes in this length */
	  			bits[j+1] += 2;		/* two new symbols in this length */
	  			bits[j]--;			/* symbol of this length is now a prefix */
	  		}
	  	}
	  	/*
	  	 * Remove the 256 symbol
	  	 */
	  	while (bits[i] == 0) {
	  		i--;
	  	}
	  	// bits[i]--;
	  	
	  	int[] finalbits = new int[16];
	  	for (i = 0; i < 16; i++) {
	  		finalbits[i] = bits[i];
	  	}
	  	return finalbits;
	}
	/**
	 * 
	 * @param bits
	 * @param codesize
	 * @return
	 */
	protected int[] makeHUFVAL(int[] bits, int[] codesize) {
		int length = 0;
		for (int i = 0; i < 16; i++) {
			length += bits[i];
		}
		int[] hufval = new int[length];
		int[] hufval_index = new int[16];
		
		hufval_index[0] = 0;
		for (int i = 1; i < 16; i++) {
			hufval_index[i] = hufval_index[i - 1] + bits[i - 1];	
		}
		
		for (int word = 0; word < codesize.length; word++) {
			int cs = codesize[word];
			if (cs > 0) {
				int index = hufval_index[cs - 1];
				hufval[index] = word;
				hufval_index[cs - 1] += 1;
			}
		}
		
		return hufval;
	}
	
	/**
	 * Per jpeglib's jpeg_gen_optimal_table method.
	 * Don't know why my coding doesn't work.
	 * @param freq holds the frequency of generated symbols
	 */
	public int[] generate_codesize_table(int[] freq) {
		int[] codesize = new int[257];		/* codesize[k] = code length of symbol k */
		int[] others = new int[257];		/* next symbol in current branch of tree */
	  	int c1 = 0, c2;
	  	
	  	for (int i = 0; i < 257; i++) {
	  		others[i] = -1;
	  	}
	  	for (int i = 0; i < 257; i++) {
	  		codesize[i] = 0;
	  	}
	  	codesize[256] = 1;
	  	boolean done = false;
	  	
	  	/*
	  	 * Making codesize
	  	 */
	  	int count = 0;
	  	while (!done) {
	  		
		  	c1 = -1;
		  	int v = Integer.MAX_VALUE;
		  	for (int i = 0; i < 257; i++) {
		  		if (freq[i] > 0 && freq[i] <= v) {
		  			v = freq[i];
		  			c1 = i;
		  		}
		  	}
		  	
		  	c2 = -1;
		  	v = Integer.MAX_VALUE;
		  	for (int i = 0; i < 257; i++) {
		  		if (freq[i] > 0 && freq[i] <= v && i != c1) {
		  			v = freq[i];
		  			c2 = i;
		  		}
		  	}
		  	
		  	if (c2 < 0) {
		  		/*
		  		 * All but c1 is zero, so it's done
		  		 */
		  		done = true;
		  	} else {
		  		count ++;
		  		freq[c1] += freq[c2];
		  		freq[c2] = 0;
		  		
		  		/*
		  		 * Increment all associated to c1's codesize 
		  		 */
		  		codesize[c1]++;
		  		while (others[c1] >= 0) {
		  			c1 = others[c1];
		  			codesize[c1]++;
		  		}
		  		
		  		/*
		  		 * c2 is now a sibling of c1
		  		 */
		  		others[c1] = c2;
		  		
		  		/*
		  		 * Increment all associated to c2's codesize
		  		 */
		  		codesize[c2]++;
		  		while (others[c2] >= 0) {
		  			c2 = others[c2];
		  			codesize[c2]++;
		  		}
		  	}
	  	}
	  	if (count == 0 && c1 >= 0) codesize[c1]++;
		return codesize;
	}
	/**
	 * Generate the CodeSize for each hufman value of this image
	 * @param FREQ Frequency of every possible word in this image
	 * @return CodeSize sequence
	 */
	public int[] generateCodeSize(int[] FREQ) {
		int length = FREQ.length;
		
		int[] clone = new int[length];
		for (int i = 0; i < length; i++) {
			clone[i] = FREQ[i];
		}
		
		int[] codesize = new int[length];
		int[] others = new int[length];
		
		/*
		 * Fills others with -1
		 */
		for (int i = 0; i < length; i++) {
			others[i] = -1;
		}
		boolean done = false;
		int count = 0;
		while (!done) {
			/*
			 * Find the smallest and 2nd smallest frequency 
			 */
			// Find the smallest
			int v1 = 0;
			for (int i = 1; i < length; i++) {
				if (clone[v1] == 0 && clone[i] > 0)
					v1 = i;
				else if (clone[i] < clone[v1] && clone[i] > 0)
					v1 = i;
			}
			// Find the 2nd smallest
			int v2 = 0;
			for (int i = 1; i < length; i++) {
				if ((clone[v2] == 0 || v2 == v1) && clone[i] > 0)
					v2 = i;
				else if (clone[i] < clone[v2] && clone[i] >= clone[v1] && i != v1)
					v2 = i;
			}
			if (clone[v2] > 0 && v2 != v1) {
				/*
				 * V2 exists.
				 * Combines 2 frequencies, they're now have the same parent.
				 * From now on, if the code size of anyone of these frequencies
				 * is incremented, the other one is incremented also. 
				 */
				clone[v1] = clone[v2] + clone[v1];
				clone[v2] = 0;

				/*
				 * Make change to v1 and every associated to v1's code size
				 */
				codesize[v1] = codesize[v1] + 1;
				while (others[v1] != -1) {
					v1 = others[v1];
					codesize[v1] = codesize[v1] + 1;
				}
				/*
				 * Now v2 and all of its associations is associated to v1
				 */
				others[v1] = v2;
				/*
				 * Make change to v2 and every associated to v2's code size
				 */
				codesize[v2] = codesize[v2] + 1;
				while (others[v2] != -1) {
					v2 = others[v2];
					codesize[v2] = codesize[v2] + 1;
				}
			} else {
				if (count == 0) codesize[v1] = 1;
				done = true;
			}
		}
		
		for (int i = 0; i < length; i++) {
			codesize[i] = (codesize[i] == 0?0:codesize[i]+1); 
		}
		return codesize;
	}
	/**
	 * Generate the frequency of every code word in this image
	 * Do the same procedure for encoding codec.huffman table but instead
	 * of encoding, counts to get frequency
	 * 
	 * @param image
	 * @return 	freg[0][0] -> lum dc
	 * 			freg[0][1] -> lum ac
	 * 			freg[1][0] -> chrom dc
	 * 			freg[1][1] -> chrom ac
	 */
	protected int[][][] generateFrequency(Image image) {		
		int[][][] freq = new int[3][2][257];
		
		int height = image.getHeight() - image.getHeight() % 8;
		int width  = image.getWidth() - image.getWidth() % 8;
		
		int[] oldDC = {0, 0, 0};
		
		Transformer t = new Transformer();
		ColorTransformer ct = new ColorTransformer();
		
		/*
		 * Convert rgb to ycc color space  
		 */
		int[][][] ycc = ct.rgb2ycc(image);

		/*
		 * Downsample image
		 */
		if (this.getHeader().isDownsample())
			ycc = ct.yccdownsample(ycc, width, height);
		for (int y = height - 1; y >= 0; y -= 8) {
			for (int x = 0; x < width; x += 8) {
				
				/*
				 * Making 8x8 mcu block
				 */
				int[][][] hexplet = new int[3][8][8];
				for (int dx = 0; dx < 8; dx++) {
					for (int dy = 0; dy < 8; dy++) {
						hexplet[0][dx][dy] = ycc[0][x+dx][y-dy];
						hexplet[1][dx][dy] = ycc[1][x+dx][y-dy];
						hexplet[2][dx][dy] = ycc[2][x+dx][y-dy];
					}
				}
				
				/*
				 * With each mcu, do a DCT transform
				 */
				double[][][] g = new double[3][8][8];
				for (int i = 0; i < 3; i++) {
					g[i] = t.transform(hexplet[i], 8, 8);
				}
				
				/*
				 * And then quantize the result transformation
				 * Which means each customized codec.huffman table is for a single
				 * quantization table only
				 */
				int[][][] q = t.quantize(g, 8, 8, this.getQuality());
				
				/*
				 * Do zigzag order with quantized data
				 */
				int[][] z = new int[3][64];
				for (int i = 0; i < 3; i++) {
					z[i] = t.zigzag(q[i], 8, 8);
				}
				
				/*
				 * Finished preparing data, start counting
				 * Those data will be used again (!) int the codec.huffman encode
				 * procedure so we should save them for later.
				 * Save the list of code word! 
				 */
				
				/*
				 * Counting
				 */
				for (int i = 0; i < 3; i++) {
					int component = (i==0)?0:1; // luminance or chrominance
					/*
					 * Count DC coef
					 */
					int[] quantized_y = z[i];
					int dc = quantized_y[0];
					int diff = dc - oldDC[i];
					oldDC[i] = dc;
					
					int size = JPEGHuffmanTable.size(diff);
					freq[component][0][size] += 1;
					
					this.getHeader().cache(diff);
					
					// Count AC coef
					int lnz = t.lastNonZero(quantized_y, 64);
					int run = 0;
					for (int j = 1; j <= lnz; j++) {
						if (quantized_y[j] == 0) {
							run++;
							if (run == 16) {
								// ZRL
								int word = 0xF0;
								run = 0;
								freq[component][1][word] += 1;
								this.getHeader().cache(word);
							}
						} else {
							int acCoef = quantized_y[j];
							size = JPEGHuffmanTable.size(acCoef);
							// (run, size)
							int word = run * 16 + size;
							run = 0;
							freq[component][1][word] += 1;
							this.getHeader().cache(word);
						}
					}
					if (quantized_y[63] == 0) {
						int word = 0x00; // eob
						freq[component][1][word] += 1;
						this.getHeader().cache(word);
					}
				}
			}
		}
		
		return freq;
	}
	
	@Override
	public int[] getEHUFCODE(int table) {
		switch(table) {
		case LUMINANCE_DC: return LUMINANCE_DC_EHUFCODE;
		case LUMINANCE_AC: return LUMINANCE_AC_EHUFCODE;
		case CHROMINANCE_DC: return CHROMINANCE_DC_EHUFCODE;
		case CHROMINANCE_AC: return CHROMINANCE_AC_EHUFCODE;
		default:
			return null;
		}
	}

	@Override
	public int[] getEHUFSIZE(int table) {
		switch(table) {
		case LUMINANCE_DC: return LUMINANCE_DC_EHUFSIZE;
		case LUMINANCE_AC: return LUMINANCE_AC_EHUFSIZE;
		case CHROMINANCE_DC: return CHROMINANCE_DC_EHUFSIZE;
		case CHROMINANCE_AC: return CHROMINANCE_AC_EHUFSIZE;
		default:
			return null;
		}
	}
	
	/**
	 * @return the image
	 */
	public Image getImage() {
		return image;
	}
	/**
	 * @param image the image to set
	 */
	public void setImage(Image image) {
		this.image = image;
	}

	/**
	 * @return the quality
	 */
	public double getQuality() {
		return this.getHeader().getQuality();
	}

	@Override
	public int[] getBITS(int component) {
		switch(component) {
		case LUMINANCE_DC: return LUMINANCE_DC_BITS;
		case LUMINANCE_AC: return LUMINANCE_AC_BITS;
		case CHROMINANCE_DC: return CHROMINANCE_DC_BITS;
		case CHROMINANCE_AC: return CHROMINANCE_AC_BITS;
		default:
			return null;
		}
	}

	@Override
	public int[] getHUFVAL(int component) {
		switch(component) {
		case LUMINANCE_DC: return LUMINANCE_DC_HUFVAL;
		case LUMINANCE_AC: return LUMINANCE_AC_HUFVAL;
		case CHROMINANCE_DC: return CHROMINANCE_DC_HUFVAL;
		case CHROMINANCE_AC: return CHROMINANCE_AC_HUFVAL;
		default:
			return null;
		}
	}

	/**
	 * @return the header
	 */
	public JPEGHeader getHeader() {
		return header;
	}

	/**
	 * @param header the header to set
	 */
	public void setHeader(JPEGHeader header) {
		this.header = header;
	}

}